package com.atguigu.transform;

import com.atguigu.func.ClickSource;
import com.atguigu.pojo.Event;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * @Author Mr.Zheng
 * @Date 2023/6/14 18:15
 *
 * 自定义分区：
 *      调用partitionCustom（Partitioner,KeySelector)
 *
 *       Partitioner:
 *          int partition(K key,int numPartitions);按照key 来计算分区号
 *              K key ：当前数据的key
 *              int numpartitions： 下游分区数
 *
 *       KeySelector：
 *          KEY getKey(IN Value) throws Exceptio; 从数据中提取key
 *
 *
 */
public class Flink06_UserDefinePartitioner {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        //自定义分区实现中外隔离
        DataStreamSource<Event> ds = env.addSource(new ClickSource());
        ds.partitionCustom(
                new Partitioner<String>() {
                    @Override
                    public int partition(String key, int numPartitions) {
                        if("Zhangs".equals(key) || "Lisi".equals(key)){
                            return 0;
                        }
                        return 1;
                    }
                }
                ,
                new KeySelector<Event, String>() {
                    @Override
                    public String getKey(Event value) throws Exception {
                        return value.getUser();
                    }
                }
        ).print().setParallelism(2);


        env.execute();
    }

}
